"""
Grad-CAM and Grad-CAM++ Visualization

实现类别激活映射可视化，包括：
- 标准Grad-CAM
- Grad-CAM++
- 多层对比
- 引导式Grad-CAM
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
import cv2
from typing import Optional, List, Tuple, Dict
import logging

logger = logging.getLogger(__name__)


class GradCAM:
    """
    Grad-CAM实现

    论文: Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization
    """

    def __init__(self, model: nn.Module, target_layer: nn.Module, device: str = 'cuda'):
        """
        Args:
            model: PyTorch模型
            target_layer: 目标层（通常是最后一个卷积层）
            device: 运行设备
        """
        self.model = model
        self.target_layer = target_layer
        self.device = device

        self.gradients = None
        self.activations = None
        self.hooks = []

        self._register_hooks()

    def _register_hooks(self):
        """注册前向和反向Hook"""

        def forward_hook(module, input, output):
            self.activations = output.detach()

        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0].detach()

        self.hooks.append(self.target_layer.register_forward_hook(forward_hook))
        self.hooks.append(self.target_layer.register_full_backward_hook(backward_hook))

    def remove_hooks(self):
        """移除Hook"""
        for hook in self.hooks:
            hook.remove()
        self.hooks = []

    def generate_cam(
        self,
        input_tensor: torch.Tensor,
        target_class: Optional[int] = None,
        roi_position: Optional[torch.Tensor] = None
    ) -> np.ndarray:
        """
        生成CAM

        Args:
            input_tensor: 输入张量 [1, 3, H, W]
            target_class: 目标类别（None则使用预测类别）
            roi_position: ROI位置信息（如果模型需要）

        Returns:
            CAM热力图 [H, W]，范围0-1
        """
        self.model.eval()
        self.gradients = None
        self.activations = None

        # 前向传播
        input_tensor.requires_grad = True

        if roi_position is not None:
            output, _, _ = self.model(input_tensor, roi_position)
        else:
            try:
                output, _, _ = self.model(input_tensor, None)
            except:
                output = self.model(input_tensor)

        # 如果output是分类logits [B, num_classes]
        if target_class is None:
            target_class = output.argmax(dim=1).item()

        # 选择目标类别的得分
        if output.dim() == 2:
            score = output[0, target_class]
        else:
            score = output[0]

        # 反向传播
        self.model.zero_grad()
        score.backward(retain_graph=True)

        # 计算权重（Global Average Pooling of gradients）
        gradients = self.gradients  # [1, C, H, W]
        activations = self.activations  # [1, C, H, W]

        if gradients is None or activations is None:
            logger.warning("Gradients or activations not captured")
            return np.zeros((input_tensor.shape[2], input_tensor.shape[3]))

        # 权重 = 梯度的全局平均
        weights = torch.mean(gradients, dim=(2, 3), keepdim=True)  # [1, C, 1, 1]

        # 加权求和
        cam = torch.sum(weights * activations, dim=1, keepdim=True)  # [1, 1, H, W]

        # ReLU（只保留正值）
        cam = F.relu(cam)

        # 上采样到输入大小
        cam = F.interpolate(
            cam,
            size=(input_tensor.shape[2], input_tensor.shape[3]),
            mode='bilinear',
            align_corners=False
        )

        # 归一化到0-1
        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)

        return cam

    def visualize_cam(
        self,
        image: np.ndarray,
        cam: np.ndarray,
        alpha: float = 0.5,
        colormap: int = cv2.COLORMAP_JET
    ) -> Image.Image:
        """
        可视化CAM叠加在原图上

        Args:
            image: 原始图像 [H, W, 3]，范围0-255
            cam: CAM热力图 [H, W]，范围0-1
            alpha: 叠加透明度
            colormap: OpenCV颜色映射

        Returns:
            PIL.Image
        """
        # CAM转为彩色热力图
        cam_uint8 = (cam * 255).astype(np.uint8)
        heatmap = cv2.applyColorMap(cam_uint8, colormap)
        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)

        # 确保尺寸匹配
        if heatmap.shape[:2] != image.shape[:2]:
            heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))

        # 叠加
        overlay = (alpha * heatmap + (1 - alpha) * image).astype(np.uint8)

        return Image.fromarray(overlay)


class GradCAMPlusPlus(GradCAM):
    """
    Grad-CAM++实现

    改进版本，对多个目标的定位更准确
    """

    def generate_cam(
        self,
        input_tensor: torch.Tensor,
        target_class: Optional[int] = None,
        roi_position: Optional[torch.Tensor] = None
    ) -> np.ndarray:
        """
        生成Grad-CAM++

        Args:
            input_tensor: 输入张量
            target_class: 目标类别
            roi_position: ROI位置

        Returns:
            CAM热力图
        """
        self.model.eval()
        self.gradients = None
        self.activations = None

        input_tensor.requires_grad = True

        if roi_position is not None:
            output, _, _ = self.model(input_tensor, roi_position)
        else:
            try:
                output, _, _ = self.model(input_tensor, None)
            except:
                output = self.model(input_tensor)

        if target_class is None:
            target_class = output.argmax(dim=1).item()

        if output.dim() == 2:
            score = output[0, target_class]
        else:
            score = output[0]

        self.model.zero_grad()
        score.backward(retain_graph=True)

        gradients = self.gradients
        activations = self.activations

        if gradients is None or activations is None:
            return np.zeros((input_tensor.shape[2], input_tensor.shape[3]))

        # Grad-CAM++权重计算
        # alpha = 梯度的二阶导数 / (2 * 二阶导数 + sum(activation * 三阶导数))

        # 简化版本：使用梯度平方作为权重
        weights = torch.mean(gradients ** 2, dim=(2, 3), keepdim=True)

        # 加权求和
        cam = torch.sum(weights * activations, dim=1, keepdim=True)
        cam = F.relu(cam)

        # 上采样
        cam = F.interpolate(
            cam,
            size=(input_tensor.shape[2], input_tensor.shape[3]),
            mode='bilinear',
            align_corners=False
        )

        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)

        return cam


class MultiLayerGradCAM:
    """多层Grad-CAM对比"""

    def __init__(self, model: nn.Module, device: str = 'cuda'):
        """
        Args:
            model: PyTorch模型
            device: 运行设备
        """
        self.model = model
        self.device = device

    def generate_multilayer_cams(
        self,
        input_tensor: torch.Tensor,
        target_layers: List[nn.Module],
        layer_names: List[str],
        target_class: Optional[int] = None,
        roi_position: Optional[torch.Tensor] = None
    ) -> Dict[str, np.ndarray]:
        """
        生成多层CAM

        Args:
            input_tensor: 输入张量
            target_layers: 目标层列表
            layer_names: 层名称列表
            target_class: 目标类别
            roi_position: ROI位置

        Returns:
            {layer_name: cam_array}
        """
        cams = {}

        for layer, name in zip(target_layers, layer_names):
            try:
                gradcam = GradCAM(self.model, layer, self.device)
                cam = gradcam.generate_cam(input_tensor, target_class, roi_position)
                cams[name] = cam
                gradcam.remove_hooks()
            except Exception as e:
                logger.error(f"Failed to generate CAM for {name}: {e}")

        return cams

    def visualize_multilayer_cams(
        self,
        image: np.ndarray,
        cams: Dict[str, np.ndarray],
        figsize: Tuple[int, int] = (20, 5)
    ) -> Image.Image:
        """
        可视化多层CAM对比

        Args:
            image: 原始图像
            cams: CAM字典
            figsize: 图像大小

        Returns:
            PIL.Image
        """
        num_layers = len(cams)
        fig, axes = plt.subplots(1, num_layers + 1, figsize=figsize)

        # 显示原图
        axes[0].imshow(image)
        axes[0].set_title('Original', fontsize=12)
        axes[0].axis('off')

        # 显示每层CAM
        for idx, (name, cam) in enumerate(cams.items(), 1):
            # 创建叠加图
            cam_uint8 = (cam * 255).astype(np.uint8)
            heatmap = cv2.applyColorMap(cam_uint8, cv2.COLORMAP_JET)
            heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)

            if heatmap.shape[:2] != image.shape[:2]:
                heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))

            overlay = (0.5 * heatmap + 0.5 * image).astype(np.uint8)

            axes[idx].imshow(overlay)
            axes[idx].set_title(name, fontsize=12)
            axes[idx].axis('off')

        plt.tight_layout()

        # 转换为PIL Image
        fig.canvas.draw()
        img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        img = Image.fromarray(img_array)

        plt.close(fig)

        return img


class GradCAMVisualizer:
    """Grad-CAM可视化统一接口"""

    def __init__(self, refiner_model, device='cuda'):
        """
        Args:
            refiner_model: Stage2 ROI Refiner模型
            device: 运行设备
        """
        self.refiner_model = refiner_model
        self.device = device

    def get_target_layers(self) -> List[Tuple[str, nn.Module]]:
        """获取可用的目标层"""
        layers = []

        # 尝试获取常见的卷积层
        if hasattr(self.refiner_model, 'fusion_module'):
            layers.append(('fusion', self.refiner_model.fusion_module))

        if hasattr(self.refiner_model, 'overlock_extractor'):
            # 可能有多个stage
            pass

        return layers

    def visualize_gradcam(
        self,
        image_tensor: torch.Tensor,
        roi_position: Optional[torch.Tensor],
        target_class: Optional[int] = None,
        target_layer_name: str = 'fusion',
        method: str = 'gradcam'
    ) -> Tuple[Image.Image, np.ndarray]:
        """
        生成Grad-CAM可视化

        Args:
            image_tensor: 输入图像张量 [1, 3, H, W]
            roi_position: ROI位置
            target_class: 目标类别
            target_layer_name: 目标层名称
            method: 'gradcam' or 'gradcam++'

        Returns:
            (可视化图像, CAM数组)
        """
        # 获取目标层
        target_layer = None
        if hasattr(self.refiner_model, target_layer_name):
            target_layer = getattr(self.refiner_model, target_layer_name)
        elif hasattr(self.refiner_model, 'fusion_module'):
            target_layer = self.refiner_model.fusion_module

        if target_layer is None:
            raise ValueError(f"Target layer '{target_layer_name}' not found")

        # 创建Grad-CAM
        if method == 'gradcam++':
            gradcam = GradCAMPlusPlus(self.refiner_model, target_layer, self.device)
        else:
            gradcam = GradCAM(self.refiner_model, target_layer, self.device)

        # 生成CAM
        cam = gradcam.generate_cam(image_tensor, target_class, roi_position)

        # 准备原图（用于叠加）
        # 反归一化
        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
        image = image_tensor[0].cpu() * std + mean
        image = (image.permute(1, 2, 0).numpy() * 255).astype(np.uint8)

        # 可视化
        vis_img = gradcam.visualize_cam(image, cam)

        # 清理
        gradcam.remove_hooks()

        return vis_img, cam


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    print("Grad-CAM Visualization Module")
    print("This module provides Grad-CAM and Grad-CAM++ visualization")
    print("Import and use in your code:")
    print("  from src.visualization.gradcam import GradCAMVisualizer")
